TIMX 533 - Load pyarrow dataset on TIMDEXDataset init by ghukill · Pull Request #160 · MITLibraries/timdex-dataset-api

ghukill · 2025-08-07T20:53:38Z

Note: this PR builds on #159.

Purpose and background context

This PR updates how a TIMDEXDataset instance is initialized.

Formerly, initialization allowed for a str (root of pyarrow dataset) or a list[str] which was an explicit list of parquet files. This reflected an earlier approach where direct manipulation of the parquet files loaded was used for filtering. Now that we're moving into an architecture where we have metadata, and the dataset has more structure (e.g. /data/records and /metadata) we want the TIMDEXDataset to represent the TIMDEX dataset as a whole.

As such, we now only allow for a location: str which is the root of the dataset. A future ticket, TIMX-529, will begin applying the metadata to queries, utilizing a different approach for limiting the parquet files.

By only allowing a str location, and no longer temporarily swapping out the dataset parquet files associated with self.dataset, a lot of complexity is removed. Somewhat projecting to TIMX-529, the approach can now be summarized as:

Init a TIMDEXDataset instance with the dataset root, e.g. s3://timdex/dataset
TIMDEXDataset uses that location and knows the pyarrrow ETL parquet dataset is at /dataset/data/records, and that dataset is automatically loaded
TIMDEXDatasetMetadata also gets the same location, knowing it can find its assets at /dataset/metadata; this is a big win for this simplification here
Eventually, read methods will apply parquet file "pruning" (e.g. only querying from a subset) based on metadata queries, making it unnecessary to "hot swap" the parquet files assocated with self.dataset

Next steps:

Rework read methods to use SQL + metadata, TIMX-529
Merge append deltas into static DB, TIMX-528

How can a reviewer manually see the effects of these changes?

1- Set AWS Dev TimdexManagers credentials

2- Set env vars:

TDA_LOG_LEVEL=DEBUG
WARNING_ONLY_LOGGERS=asyncio,botocore,urllib3,s3transfer,boto3,MARKDOWN
TIMDEX_DATASET_LOCATION=s3://timdex-extract-dev-222053980223/dataset_scratch

3- Start Ipython shell with pipenv run ipython and do some setup:

import os

from timdex_dataset_api import TIMDEXDataset
from timdex_dataset_api.config import configure_dev_logger
from tests.utils import generate_sample_records

configure_dev_logger()

4- Create a TIMDEXDataset instance, noting that a pyarrow dataset is automatically loaded:

td = TIMDEXDataset(os.environ["TIMDEX_DATASET_LOCATION"])
"""
...
INFO:timdex_dataset_api.dataset:Dataset successfully loaded: 's3://timdex-extract-dev-222053980223/dataset_scratch/data/records', 1.06s
...
"""

5- Note the removal of .load(), replaced by a more explicit .load_pyarrow_dataset() which returns a dataset instance versus setting on self:

pa_ds = td.load_pyarrow_dataset()
"""
INFO:timdex_dataset_api.dataset:Dataset successfully loaded: 's3://timdex-extract-dev-222053980223/dataset_scratch/data/records', 1.06s
"""

type(pa_ds)
"""
Out[3]: pyarrow._dataset.FileSystemDataset
"""

Includes new or updated dependencies?

NO

Changes expectations for external applications?

NO

What are the relevant tickets?

https://mitlibraries.atlassian.net/browse/TIMX-533

coveralls · 2025-08-07T20:56:31Z

Pull Request Test Coverage Report for Build 16840527470

Details

38 of 38 (100.0%) changed or added relevant lines in 1 file are covered.
1 unchanged line in 1 file lost coverage.
Overall coverage increased (+0.3%) to 94.828%

Files with Coverage Reduction	New Missed Lines	%
timdex_dataset_api/exceptions.py	1	0.0%

Totals
Change from base Build 16840506568:	0.3%
Covered Lines:	385
Relevant Lines:	406

💛 - Coveralls

ghukill · 2025-08-08T13:27:19Z

-            location (str | list[str]): Local filesystem path or an S3 URI to
-                a parquet dataset. For partitioned datasets, set to the base directory.
+            location (str ): Local filesystem path or an S3 URI to a parquet dataset.


This is a big change, that reverberates throughout. By limiting to only str here, we're forced (get to!) remove the list[str] options anywhere self.location was used.

ghukill · 2025-08-08T13:32:04Z

+        # set source for pyarrow dataset
+        source: str | list[str] = parquet_files or path


Unlike TIMDEXDataset where we only ever want the root of the whole dataset, including metadata, a pyarrow parquet dataset is happy to accept a str or list[str].

There may be instances where we use TIMDEXDataset.load_pyarrow_dataset() to load a specific, temporary dataset from a list of parquet files, which is why this method allows the optional arg parquet_files: list[str] | None.

ghukill · 2025-08-08T13:32:40Z

-    def _get_filtered_dataset(
-        self,
-        **filters: Unpack[DatasetFilters],
-    ) -> ds.Dataset:


This method isn't removed, just moved! But.... it may get removed in TIMX-529.

ghukill · 2025-08-08T13:34:00Z

+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_all_current_records_deduped(timdex_dataset_with_runs):
+    timdex_dataset_with_runs.load(current_records=True)
+    all_records_df = timdex_dataset_with_runs.read_dataframe()
+
+    # assert both sources have accurate record counts for current records only
+    assert all_records_df.source.value_counts().to_dict() == {"dspace": 90, "alma": 100}
+
+    # assert only one "full" run, per source
+    assert len(all_records_df[all_records_df.run_type == "full"].run_id.unique()) == 2
+
+    # assert run_date min/max dates align with both sources min/max dates
+    assert all_records_df.run_date.min() == date(2025, 1, 1)  # both
+    assert all_records_df.run_date.max() == date(2025, 2, 5)  # dspace
+
+
+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_source_current_records_deduped(timdex_dataset_with_runs):
+    timdex_dataset_with_runs.load(current_records=True, source="alma")
+    alma_records_df = timdex_dataset_with_runs.read_dataframe()
+
+    # assert only alma records present and correct count
+    assert alma_records_df.source.value_counts().to_dict() == {"alma": 100}
+
+    # assert only one "full" run
+    assert len(alma_records_df[alma_records_df.run_type == "full"].run_id.unique()) == 1
+
+    # assert run_date min/max dates are correct for single source
+    assert alma_records_df.run_date.min() == date(2025, 1, 1)
+    assert alma_records_df.run_date.max() == date(2025, 1, 5)
+
+
+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_all_read_methods_get_deduplication(
+    timdex_dataset_with_runs,
+):
+    timdex_dataset_with_runs.load(current_records=True, source="alma")
+
+    full_df = timdex_dataset_with_runs.read_dataframe()
+    all_records = list(timdex_dataset_with_runs.read_dicts_iter())
+    transformed_records = list(timdex_dataset_with_runs.read_transformed_records_iter())
+
+    assert len(full_df) == len(all_records) == len(transformed_records)
+
+
+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_current_records_no_additional_filtering_accurate_records_yielded(
+    timdex_dataset_with_runs,
+):
+    timdex_dataset_with_runs.load(current_records=True, source="alma")
+    df = timdex_dataset_with_runs.read_dataframe()
+    assert df.action.value_counts().to_dict() == {"index": 99, "delete": 1}
+
+
+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_current_records_action_filtering_accurate_records_yielded(
+    timdex_dataset_with_runs,
+):
+    timdex_dataset_with_runs.load(current_records=True, source="alma")
+    df = timdex_dataset_with_runs.read_dataframe(action="index")
+    assert df.action.value_counts().to_dict() == {"index": 99}
+
+
+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_current_records_index_filtering_accurate_records_yielded(
+    timdex_dataset_with_runs,
+):
+    """This is a somewhat complex test, but demonstrates that only 'current' records
+    are yielded when .load(current_records=True) is applied.
+
+    Given these runs from the fixture:
+    [
+        ...
+        (25, "alma", "2025-01-03", "daily", "index", "run-5"),   <---- filtered to
+        (10, "alma", "2025-01-04", "daily", "delete", "run-6"),  <---- influences current
+        ...
+    ]
+
+    Though we are filtering to run-5, which has 25 total records to-index, we see only 15
+    records yielded.  Why?  This is because while we have filtered to only yield from
+    run-5, run-6 had 10 deletes which made records alma:0|9 no longer "current" in run-5.
+    As we yielded records reverse chronologically, the deletes from run-6 (alma:0-alma:9)
+    "influenced" what records we would see as we continue backwards in time.
+    """
+    # with current_records=False, we get all 25 records from run-5
+    timdex_dataset_with_runs.load(current_records=False, source="alma")
+    df = timdex_dataset_with_runs.read_dataframe(run_id="run-5")
+    assert len(df) == 25
+
+    # with current_records=True, we only get 15 records from run-5
+    # because newer run-6 influenced what records are current for older run-5
+    timdex_dataset_with_runs.load(current_records=True, source="alma")
+    df = timdex_dataset_with_runs.read_dataframe(run_id="run-5")
+    assert len(df) == 15
+    assert list(df.timdex_record_id) == [
+        "alma:10",
+        "alma:11",
+        "alma:12",
+        "alma:13",
+        "alma:14",
+        "alma:15",
+        "alma:16",
+        "alma:17",
+        "alma:18",
+        "alma:19",
+        "alma:20",
+        "alma:21",
+        "alma:22",
+        "alma:23",
+        "alma:24",
+    ]
+
+
+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_load_current_records_gets_correct_same_day_full_run(
+    timdex_dataset_same_day_runs,
+):
+    """Two full runs were performed on the same day, but 'run-2' was performed most
+    recently.  current_records=True should discover the more recent of the two 'run-2',
+    not 'run-1'."""
+    timdex_dataset_same_day_runs.load(current_records=True, run_type="full")
+    df = timdex_dataset_same_day_runs.read_dataframe()
+
+    assert list(df.run_id.unique()) == ["run-2"]
+
+
+@pytest.mark.skip(reason="All tests for 'current' records will be reworked.")
+def test_dataset_load_current_records_gets_correct_same_day_daily_runs_ordering(
+    timdex_dataset_same_day_runs,
+):
+    """Two runs were performed on 2025-01-02, but the most recent records should be from
+    run 'run-5' which are action='delete', not 'run-4' with action='index'."""
+    timdex_dataset_same_day_runs.load(current_records=True, run_type="daily")
+    first_record = next(timdex_dataset_same_day_runs.read_dicts_iter())
+
+    assert first_record["run_id"] == "run-5"
+    assert first_record["action"] == "delete"


These were moved from tests/test_dataset.py, where they probably should have been in the test_read.py file this whole time.

They will get reworked in TIMX-529, maybe significantly, but what they test -- mostly around current records -- will remain the same.

ghukill · 2025-08-08T13:34:47Z

-    @property
-    def row_count(self) -> int:
-        """Get row count from loaded dataset."""
-        if not self.dataset:
-            raise DatasetNotLoadedError
-        return self.dataset.count_rows()


This was from like the first commit, and we really just don't need this misdirection. Better to use self.dataset.count_rows(). Decision was made easier now that we don't have a possible DatasetNotLoadedError situation.

Why these changes are being introduced: As the TIMDEXDatasetMetadata becomes more integrated, there is less need to be explicit about how we load the pyarrow dataset. Formerly, the method .load() needed to be called manually and supported options like 'current_records' or 'include_parquet_files'. This also reflected a time when 'TIMDEXDataset.load()' suggested that "loading" was the pyarrow dataset only. With the introduction of metadata, it is also better to be specific we are loading a pyarrow dataset which is only one of many assets associated with a TIMDEXDataset instance. How this addresses that need: Renames .load() to .load_pyarrow_dataset() to be explicit about what is happening. We no longer store the pyarrow dataset filesystem or paths on self, as they are only used briefly during this dataset load. We can get them anytime via .dataset. Really most important, we limit the root 'location' that we init a TIMDEXDataset instance to be a string only, the root of the dataset. Now that we don't allow a list of strings at that level, we can trust the nature of self.location to be a string, and the root of the TIMDEX dataset. Side effects of this change: * TIMDEXDataset and TIMDEXDatasetMetadata can only be initialized with a string, which is the root of the TIMDEX dataset. From there, both know where their assets can be found. * You cannot "pre-filter" the pyarrow dataset when loading, which had confusing overlap with the read methods; the read methods themselves may change somewhat dramatically now that we have metadata to use. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-533

ehanson8

Always appreciate when the red number is bigger than the green number! ipython test worked as expected, I like the explicitness that's developed over these commits. Great work!

ghukill · 2025-08-11T17:54:38Z

Always appreciate when the red number is bigger than the green number!

I lol'ed! And agreed. I think TIMX-529 may be about even, but it may tip red too...

ghukill force-pushed the TIMX-533-rework-dataset-load branch from 4412f0f to d9df13b Compare August 7, 2025 20:58

ghukill requested a review from a team August 8, 2025 13:25

ghukill commented Aug 8, 2025

View reviewed changes

ghukill force-pushed the TIMX-526-projected-views branch from 05383bc to 43e5350 Compare August 8, 2025 17:48

ghukill changed the base branch from TIMX-526-projected-views to epic-TIMX-515 August 8, 2025 21:05

ghukill force-pushed the TIMX-533-rework-dataset-load branch from d9df13b to 6f75254 Compare August 8, 2025 21:05

ghukill marked this pull request as ready for review August 8, 2025 21:06

ehanson8 self-assigned this Aug 11, 2025

ehanson8 approved these changes Aug 11, 2025

View reviewed changes

ghukill merged commit e63d29d into epic-TIMX-515 Aug 11, 2025
2 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

TIMX 533 - Load pyarrow dataset on TIMDEXDataset init#160

TIMX 533 - Load pyarrow dataset on TIMDEXDataset init#160
ghukill merged 1 commit into
epic-TIMX-515from
TIMX-533-rework-dataset-load

ghukill commented Aug 7, 2025 •

edited

Loading

Uh oh!

coveralls commented Aug 7, 2025 •

edited

Loading

Uh oh!

ghukill Aug 8, 2025

Uh oh!

ghukill Aug 8, 2025

Uh oh!

ghukill Aug 8, 2025

Uh oh!

ghukill Aug 8, 2025

Uh oh!

ghukill Aug 8, 2025 •

edited

Loading

Uh oh!

ehanson8 left a comment

Uh oh!

ghukill commented Aug 11, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		# set source for pyarrow dataset
		source: str \| list[str] = parquet_files or path

Uh oh!

Conversation

ghukill commented Aug 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose and background context

Next steps:

How can a reviewer manually see the effects of these changes?

Includes new or updated dependencies?

Changes expectations for external applications?

What are the relevant tickets?

Uh oh!

coveralls commented Aug 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 16840527470

Details

💛 - Coveralls

Uh oh!

ghukill Aug 8, 2025

Choose a reason for hiding this comment

Uh oh!

ghukill Aug 8, 2025

Choose a reason for hiding this comment

Uh oh!

ghukill Aug 8, 2025

Choose a reason for hiding this comment

Uh oh!

ghukill Aug 8, 2025

Choose a reason for hiding this comment

Uh oh!

ghukill Aug 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ehanson8 left a comment

Choose a reason for hiding this comment

Uh oh!

ghukill commented Aug 11, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ghukill commented Aug 7, 2025 •

edited

Loading

coveralls commented Aug 7, 2025 •

edited

Loading

ghukill Aug 8, 2025 •

edited

Loading